from matplotlib import collections
import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.axisartist as axisartist

def sa(t, shift, use_bias):
    
    seed = np.abs(int(t * 10))
    np.random.seed(seed)
    if use_bias:
        bias = np.random.uniform(-0.05, 0.05)
    else:
        bias = 0
    t -= shift
    
    if t == 0:
        return 1
    return np.sin(t) / t + bias
    
def set_axis(fig, number, title, x_range=(-10, 15), y_range=(-1, 1)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(x_range[0], x_range[1])
    ax.set_ylim(y_range[0], y_range[1])
    ax.set_title(title)
    
def correlate(tau, t, fun, shift):
    return np.sum(np.array([fun(i, shift, True) * fun(i - tau, 0, False) for i in t]))


if __name__ == "__main__":
    
    fig = plt.figure(figsize=(12, 4))
    plt.rcParams['font.family'] = ['Arial Unicode MS']
    
    limit = 8
    delta = 2
    t = np.arange(-3 * limit, 3 * limit, 0.1)

    fun1 = np.array([sa(i, 2, True) for i in t])
    set_axis(fig, 131, "x(t)", (-3 * limit, 3 * limit), (-1, 2))
    plt.plot(t, fun1)
    
    fun2 = np.array([sa(i, 0, False) for i in t])
    set_axis(fig, 132, "s(t)", (-3 * limit, 3 * limit), (-1, 2))
    plt.plot(t, fun2)
    
    fun3 = np.array([correlate(i, t, sa, 2) for i in t])
    set_axis(fig, 133, "$r_{xs}$", (-3 * limit, 3 * limit), (-10, 35))
    # print(fun3)
    plt.plot(t, fun3)
    
    plt.show()